'''
https://leetcode.cn/problems/falling-squares/description/
'''
from typing import List

import numpy as np


class SegmentTree:
    def __init__(self, n):
        self.n = n
        self.info = [0] * (n << 2)
        self.lazy_update = [False] * (n << 2)
        self.change = [0] * (n << 2)

    def _up(self, root):
        self.info[root] = max(self.info[root << 1], self.info[root << 1 | 1])

    def _down(self, root):
        if self.lazy_update[root]:
            self._lazy(root << 1, self.change[root])
            self._lazy(root << 1 | 1, self.change[root])
            self.lazy_update[root] = False

    def _lazy(self, root, v):
        self.info[root] = v
        self.lazy_update[root] = True
        self.change[root] = v

    def update(self, L, R, V):
        def f(l, r, root):
            if r < L or l > R: return
            if L <= l <= r <= R:
                self._lazy(root, V)
                return
            self._down(root)
            mid = (l + r) >> 1
            f(l, mid, root << 1)
            f(mid + 1, r, root << 1 | 1)
            self._up(root)

        f(1, self.n, 1)

    def query(self, L, R):
        def f(l, r, root):
            # same case like add
            if r < L or l > R: return 0
            if L <= l <= r <= R:
                return self.info[root]
            self._down(root)
            mid = (l + r) >> 1
            return max(f(l, mid, root << 1), f(mid + 1, r, root << 1 | 1))

        return f(1, self.n, 1)


class Solution:
    def fallingSquares(self, positions: List[List[int]]) -> List[int]:
        dict = {}

        def discret():
            all_point = set()
            for x, side_len in positions:
                all_point.add(x)
                all_point.add(x + side_len)
            all_point = sorted(all_point)
            id = 1
            for x in all_point:
                dict[x] = id
                id += 1
            return len(all_point)

        n = discret()
        tree = SegmentTree(n)
        res = [0] * len(positions)
        history_max = 0
        for i, position in enumerate(positions):
            x, side_len = position

            range_height = tree.query(dict[x], dict[x + side_len] - 1) + side_len
            tree.update(dict[x], dict[x + side_len] - 1, range_height)

            history_max = max(range_height, history_max)
            res[i] = history_max
        return res

    # 离散化+暴力解
    def fallingSquares2(self, positions: List[List[int]]) -> List[int]:
        dict = {}

        def discrete():
            all_point = set()
            for x, side_len in positions:
                all_point.add(x)
                all_point.add(x + side_len)
            all_point = sorted(all_point)
            id = 0
            for x in all_point:
                dict[x] = id
                id += 1
            return id

        n = discrete()
        info = [0] * n
        history_max = 0
        res = [0] * len(positions)
        for i, position in enumerate(positions):
            x, side_len = position
            left, right = dict[x], dict[x + side_len]   # 不包含right

            range_height = max(info[left: right]) + side_len
            for t in range(left, right):
                info[t] = range_height
            history_max = max(range_height, history_max)
            res[i] = history_max
        return res

    # 离散化+暴力解, 试试numpy的区间修改
    import numpy as np
    def fallingSquares3(self, positions: List[List[int]]) -> List[int]:
        dict = {}

        def discrete():
            all_point = set()
            for x, side_len in positions:
                all_point.add(x)
                all_point.add(x + side_len)
            all_point = sorted(all_point)
            id = 0
            for x in all_point:
                dict[x] = id
                id += 1
            return id

        n = discrete()
        info = np.zeros(n,dtype=int)
        history_max = 0
        res = [0] * len(positions)
        for i, position in enumerate(positions):
            x, side_len = position
            left, right = dict[x], dict[x + side_len]   # 不包含right

            range_height = max(info[left: right]) + side_len
            info[left:right] = range_height
            history_max = max(range_height, history_max)
            res[i] = int(history_max)
        return res


positions = [[2, 1], [2, 9], [1, 8]]
print(Solution().fallingSquares3(positions))
